import copy
import socket
import os
import sys
import struct
import threading
import torch
from LeNet_model_MNIST import LeNet_edge_side, LeNet_client_side, LeNet_Global_Model


import asyncio
import websockets
import json
mock = [11.45,22.78,56.99, 89.999, 90.7]
# 回应前端请求
async def handle_client(websocket, path):
    client_ip = websocket.remote_address[0]
    print(f"client connected from {client_ip}")

    try:
        async for message in websocket:
            if message == 'get':
                print(client_ip + " request get")
                # 当接收到 'get' 消息时，发送 ar 变量的值
                # await websocket.send(json.dumps({'ack': 1,'accuracy': accuracy_per_epoch['5558']}))
                await websocket.send(json.dumps({'ack': 1,'accuracy': mock}))

    except websockets.exceptions.ConnectionClosed:
        pass

def run_server():
    asyncio.set_event_loop(asyncio.new_event_loop())
    start_server = websockets.serve(handle_client, '0.0.0.0', 5600)
    print('websocket run at localhost:5600')
    asyncio.get_event_loop().run_until_complete(start_server)
    asyncio.get_event_loop().run_forever()

model_kind = '0'  # 模型种类
connected = 0  # 连接个数
model = 0  # 接收模型个数
if_train = 0  # 是否开始训练
lock_con = threading.Lock()
lock_model = threading.Lock()
lock_train = threading.Lock()
lock_accuracy = threading.Lock()
is_set = 0  # 前端设置完毕
train_over = 0
accuracy_over = 0
accuracy = []
average_accuracy = 0.0
accuracy_per_epoch = {}
final_client_path = './final_client.pth'
final_server_path = './final_server.pth'


def send_string(con, string):  # 向客户端发送关闭指令
    con.send(string.encode())


def send_file(con, file_addr):
    while True:
        filepath = file_addr  # 传输文件的路径
        if os.path.isfile(filepath):
            # 定义定义文件信息。128s表示文件名为128bytes长，l表示一个int或log文件类型，在此为文件大小
            fileinfo_size = struct.calcsize('128sq')
            # 定义文件头信息，包含文件名和文件大小
            fhead = struct.pack('128sq', bytes(os.path.basename(filepath).encode('utf-8')), os.stat(filepath).st_size)
            con.send(fhead)
            print('client filepath: {0}'.format(filepath))
            with open(filepath, 'rb') as fp:
                while True:
                    data = fp.read(1024)
                    if not data:
                        print('{0} file send over...'.format(filepath))
                        break
                    con.send(data)
        break


def socket_service(port):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        # 修改ip,此处ip必须为服务器端的ip ,linux做服务器输入ifconfig得到ip
        s.bind(('0.0.0.0', port))
        s.listen(10)

    except socket.error as msg:
        print(msg)
        sys.exit(1)
    print("Waiting...")
    while True:
        conn, addr = s.accept()
        print('successful')
        t = threading.Thread(target=cloud_server, args=(conn, str(port)))
        t.start()
        break


def cloud_server(conn, port):
    global connected, if_train, accuracy_over

    command = conn.recv(1024).decode()

    if command == 'ready':
        lock_con.acquire()
        connected += 1
        print('ready', connected)
        lock_con.release()

    while True:
        if connected == 2:
            break

    command = conn.recv(1024).decode()

    if command == 'all-ready' + port:
        lock_train.acquire()
        if_train += 1
        print('all-ready' + port)
        lock_train.release()

    while True:
        if if_train == 2:
            send_string(conn, 'train')
            break

    accuracy_per_epoch[port] = []

    while True:
        command = conn.recv(1024).decode()
        send_string(conn, 'recvd')
        if command == 'over':
            break
        else:
            lock_accuracy.acquire()
            accuracy_per_epoch[port].append(command)
            lock_accuracy.release()

    recv_file(conn, port)
    send_string(conn, 'recvd')
    recv_file(conn, port)

    while True:
        if train_over == 1:
            send_file(conn, './final_client.pth')
            while True:
                command = conn.recv(1024).decode()
                if command == 'recvd':
                    break
            send_file(conn, './final_server.pth')
            break

    while True:
        command = conn.recv(1024).decode()
        send_string(conn, 'recvd')
        if command == 'over':
            lock_model.acquire()
            accuracy_over += 1
            lock_model.release()
            break
        else:
            lock_accuracy.acquire()
            accuracy.append(command)
            lock_accuracy.release()


def recv_file(conn, port):
    while True:
        fileinfo_size = struct.calcsize('128sq')  # linux 和 windows 互传 128sl 改为 128sq  机器位数不一样，一个32位一个64位
        buf = conn.recv(fileinfo_size)
        print('收到的字节流：', buf, type(buf))
        if buf:
            print(buf, type(buf))
            filename, filesize = struct.unpack('128sq', buf)
            fn = filename.strip(str.encode('\00'))
            new_filename = os.path.join(str.encode('./'), str.encode(port) + fn)
            print('file new name is {0}, filesize if {1}'.format(new_filename, filesize))
            recvd_size = 0  # 定义已接收文件的大小
            with open(new_filename, 'wb') as fp:
                print("start receiving...")
                while not recvd_size == filesize:
                    if filesize - recvd_size > 1024:
                        data = conn.recv(1024)
                        recvd_size += len(data)
                    else:
                        data = conn.recv(filesize - recvd_size)
                        recvd_size = filesize

                    fp.write(data)
            print("end receive...")
            return new_filename.decode()
        break


def FedAvg(w):
    w_avg = copy.deepcopy(w[0])
    for k in w_avg.keys():
        for i in range(1, len(w)):
            w_avg[k] += w[i][k]
        w_avg[k] = torch.div(w_avg[k], len(w))
    return w_avg


if __name__ == '__main__':
    # websockets
    server_thread = threading.Thread(target=run_server)
    server_thread.start()

    p1 = 5558
    p2 = 5559
    t1 = threading.Thread(target=socket_service, args=[5558])
    t2 = threading.Thread(target=socket_service, args=[5559])
    t1.start()
    t2.start()
    '''
    socket_service(5558)
    socket_service(5559)
    '''
    while True:
        if os.path.exists('./5558client.pth') and os.path.exists('./5559client.pth'):
            '''
            聚合算法
            假设结果为client
            '''
            net1 = LeNet_client_side()
            net2 = LeNet_client_side()
            net1.load_state_dict(torch.load('./5558client.pth'))
            net2.load_state_dict(torch.load('./5559client.pth'))
            w = [net1.state_dict(), net2.state_dict()]

            final = FedAvg(w)
            torch.save(final, './final_client.pth')
            break

    while True:
        if os.path.exists('./5558server.pth') and os.path.exists('./5559server.pth'):
            '''
            聚合算法
            假设结果为server
            聚合完将两个模型整合为一个模型
            交给用户
            
            后续处理（如有必要）
            '''
            net1 = LeNet_edge_side()
            net2 = LeNet_edge_side()
            net1.load_state_dict(torch.load('./5558server.pth'))
            net2.load_state_dict(torch.load('./5559server.pth'))
            w = [net1.state_dict(), net2.state_dict()]

            final = FedAvg(w)

            torch.save(final, './final_server.pth')
            break

    train_over = 1

    '''
    最终模型Model完成，发送回终端检测准确率 
    假设上述整合完成后模型模型存储地址为./Model.pth，即
    
    torch.save(Model.stat_dict(), ./Model.pth)
    '''
    cloud = LeNet_Global_Model()
    cloud.load_model(final_client_path, final_server_path)

    torch.save(cloud.state_dict(), './user_model.pth')

    while True:
        if accuracy_over == 2:
            for a in accuracy:
                average_accuracy += float(a)
            average_accuracy = average_accuracy / len(accuracy)
            break

    # 拼接client和server为一个模型,将模型和平均准确率返回给用户
    print(accuracy_per_epoch.keys())
    print(accuracy_per_epoch.values())